from typing import Callable, Iterable, Iterator, Self

import numpy as np


class ChannelsTilesIterator(Iterator[tuple[np.ndarray, np.ndarray, np.ndarray]]):

    @classmethod
    def create(cls, *,
        docs: Iterable[tuple[np.ndarray, int]],
        doc_tiler: Callable[[np.ndarray, int], Iterator[tuple[np.ndarray, int]]],
        channels_n: int,
    ) -> Self:
        return cls(docs, doc_tiler, channels_n)

    @classmethod
    def resume(cls, *,
        docs: Iterable[tuple[np.ndarray, int]],
        doc_tiler: Callable[[np.ndarray, int], Iterator[tuple[np.ndarray, int]]],
        channels_n: int,
        channels_doc_index: Iterable[int],
        channels_tile_offset: Iterable[int],
    ) -> Self:
        docs_iterator = iter(docs)
        channels_doc_index = list(channels_doc_index)
        channels_tile_offset = list(channels_tile_offset)

        channels_tiles_iterator: list[Iterator[tuple[np.ndarray, int]] | None] = [None] * len(channels_doc_index)
        for channel_i in range(len(channels_doc_index)):
            while channels_tiles_iterator[channel_i] is None:
                doc_tokens, doc_index = next(docs_iterator)
                for channel_j in range(channel_i, len(channels_doc_index)):
                    if doc_index == channels_doc_index[channel_j]:
                        tiles_iterator = doc_tiler(doc_tokens, channels_tile_offset[channel_j])
                        next(tiles_iterator)  # skip the iterated tile
                        channels_tiles_iterator[channel_j] = tiles_iterator
                        break

        return cls(docs_iterator, doc_tiler, channels_n, channels_doc_index, channels_tiles_iterator)

    def __init__(self,
        docs: Iterable[tuple[np.ndarray, int]],
        doc_tiler: Callable[[np.ndarray, int], Iterator[tuple[np.ndarray, int]]],
        channels_n: int,
        channels_doc_i: Iterable[int] = (),
        channels_tiles_iterator: Iterable[Iterator[tuple[np.ndarray, int]]] = (),
    ):
        self._docs_iterator = iter(docs)
        self._doc_tiler = doc_tiler
        self._channels_n = channels_n
        self._channels_doc_index = list(channels_doc_i)
        self._channels_tiles_iterator = list(channels_tiles_iterator)

        if not len(self._channels_doc_index) == len(self._channels_tiles_iterator):
            raise ValueError(
                f"Expected len(channels_doc_index)==len(channels_tiles_iterator), "
                f"but len(channels_doc_index)={len(self._channels_doc_index)} "
                f"and len(channels_tiles_iterator)={len(self._channels_tiles_iterator)}")

    def __next__(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        while len(self._channels_doc_index) < self._channels_n:
            doc_tokens, doc_index = next(self._docs_iterator)
            self._channels_doc_index.append(doc_index)
            self._channels_tiles_iterator.append(self._doc_tiler(doc_tokens, 0))

        batch_tile_tokens = []
        batch_tile_offset = []
        for channel_i in range(self._channels_n):
            while True:
                try:
                    tile_tokens, tile_offset = next(self._channels_tiles_iterator[channel_i])
                    batch_tile_tokens.append(tile_tokens)
                    batch_tile_offset.append(tile_offset)
                    break
                except StopIteration:
                    if len(self._channels_doc_index) > self._channels_n:
                        self._channels_doc_index[channel_i] = self._channels_doc_index.pop()
                        self._channels_tiles_iterator[channel_i] = self._channels_tiles_iterator.pop()
                        print(self._channels_doc_index)
                        print(self._channels_tiles_iterator)
                    else:
                        doc_tokens, doc_index = next(self._docs_iterator)
                        self._channels_doc_index[channel_i] = doc_index
                        self._channels_tiles_iterator[channel_i] = self._doc_tiler(doc_tokens, 0)
                    continue

        batch_tile_tokens = np.stack(batch_tile_tokens, axis=0)
        batch_tile_offset = np.asarray(batch_tile_offset, dtype=np.int64)
        batch_doc_index = np.asarray(self._channels_doc_index, dtype=np.int64)
        return batch_tile_tokens, batch_tile_offset, batch_doc_index
